{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Custom Dataset  \n",
    "\n",
    " <div style=\"text-align: right\"> by <a href=\"https://scholar.google.com/citations?user=f-4YHeMAAAAJ&hl=en\">Lorenzo Bertoni</a>  13/04/2021 </div>\n",
    "\n",
    "\n",
    "\n",
    "## Overview\n",
    "In this section of the guide, we will see how to train and evaluate OpenPifPaf on a custom dataset. OpenPifPaf is based on the concept of __[Plugin architecture pattern](https://cs.uwaterloo.ca/~m2nagapp/courses/CS446/1195/Arch_Design_Activity/PlugIn.pdf)__, and the overall system is composed of a core component and auxiliary plug-in modules. To train a model on a custom dataset, you don't need to change the core system, only to create a small plugin for it.\n",
    "This tutorial will go through the steps required to create a new plugin for a custom dataset.\n",
    "\n",
    "\n",
    "Let's go through the steps of implementing a 2D pose estimator for vehicles, as a case study. If you are interested in how this specific plugin works, please check its {doc}`guide section <plugins_apollocar3d>`. We suggest to create and debug your own plugin copying a pre-existing plugin, changing its name, and adapting its files to your needs.\n",
    "Below, a description of the structure of the plugin to give you some intuition of what you will need to change.\n",
    "\n",
    "##  Plugin structure\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. Data Module\n",
    "\n",
    "This module handles the interface of your custom dataset with the core system and it is the main component of the plugin. For the [ApolloCar3D Dataset](http://apolloscape.auto/car_instance.html), we created a module called _apollo_kp.py_ containing the class [ApolloKp](https://github.com/openpifpaf/openpifpaf/blob/main/openpifpaf/plugins/apollocar3d/apollo_kp.py) which inherits from the [DataModule](https://github.com/openpifpaf/openpifpaf/blob/main/openpifpaf/datasets/module.py) class.\n",
    "\n",
    "The base class to inherit from has the following structure:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{eval-rst}\n",
    ".. autoclass:: openpifpaf.datasets.DataModule\n",
    "    :members: cli, configure, train_loader, val_loader, eval_loader, metrics\n",
    "    :noindex:\n",
    "```\n",
    "\n",
    "Now that you have a general view of the structure of a data module, we suggest you to refer to the implementation of the [ApolloKp](https://github.com/openpifpaf/openpifpaf/blob/main/openpifpaf/plugins/apollocar3d/apollo_kp.py) class. You can get started by copying and modifying this class according to your needs."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. Plugin Registration"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For the core system to recognize the new plugin you need to create a [\\_\\_init\\_\\_.py](https://github.com/openpifpaf/openpifpaf/blob/main/openpifpaf/plugins/apollocar3d/__init__.py) file specifying:\n",
    "\n",
    "1. __Name Convention__: include all the plugins file into a folder named __openpifpaf\\_\\<plugin_name\\>__. Only folders, which names start with openpifpaf\\_ are recognized\n",
    "2. __Registration__: Inside the folder, create an _init.py_ file and add to the list of existing plugins the datamodule that we have just created (__ApolloKp__). In this case, the name __apollo__ represents the name of the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def register():\n",
    "    openpifpaf.DATAMODULES['apollo'] = ApolloKp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3. Constants\n",
    "\n",
    "Create a module _constants.py_ containing all the constants needed to define the 2D keypoints of vehicles. The most important are:\n",
    "-  __Names__ of the keypoints, as a list of strings. In our plugin this is called CAR_KEYPOINTS\n",
    "- __Skeleton__: the connections between the keypooints, as a list of lists of two elements indicating the indeces of the starting and ending connections. In our plugin this is called CAR_SKELETON\n",
    "- __Sigmas__: the size of the area to compute the object keypoint similarity (OKS), if you wish to use average precision (AP) as a metric.\n",
    "- __Score weights__:the weights to compute the overall score of an object (e.g. car or person). When computing the overall score the highest weights will be assigned to the most confident joints.\n",
    "- __Categories__ of the keypoints. In this case, the only category is car.\n",
    "- __Standard pose__ of the keypoints, to visualize the connections between the keypoints and as an argument for the head network.\n",
    "- __Horizontal flip__ equivalents, if you use horizontal flipping as augmentation technique, you will need to define the corresponding left and right and keypoints as a dictionary. E.g. left_ear --> right_ear.  In our plugin this is called HFLIP\n",
    "\n",
    "In addition to the constants, the module contains two functions to draw the skeleton and save it as an image. The functions are only for debugging and can usually be used as they are, only changing the arguments with the new constants. For additional information, refer to the file [constants.py](https://github.com/openpifpaf/openpifpaf/blob/main/openpifpaf/plugins/apollocar3d/constants.py)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4. Data Loading\n",
    "\n",
    "If  you are using COCO-style annotations, there is no need to create a datalader to load images and annotations. A default [CocoLoader](https://github.com/openpifpaf/openpifpaf/blob/main/openpifpaf/plugins/coco/dataset.py) is already available to be called inside the data module [ApolloKp](https://github.com/openpifpaf/openpifpaf/blob/main/openpifpaf/plugins/apollocar3d/apollo_kp.py)\n",
    "\n",
    "If you wish to load different annotations, you can either write your own dataloader, or you can transform your annotations to COCO style .json files. In this plugin, we first convert ApolloCar3D annotations into COCO style .json files and then load them as standard annotations."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 5. Annotations format \n",
    "This step transforms custom annotations, in this case from ApolloCar3D, into COCO-style annotations. Below we describe how to populate a json file using COCO-style format. For the full working example, check the module [apollo_to_coco.py](https://github.com/openpifpaf/openpifpaf/blob/main/openpifpaf/plugins/apollocar3d/apollo_to_coco.py) inside the plugin."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def initiate_json(self):\n",
    "    \"\"\"\n",
    "    Initiate json file: one for training phase and another one for validation.\n",
    "    \"\"\"\n",
    "    self.json_file[\"info\"] = dict(url=\"https://github.com/openpifpaf/openpifpaf\",\n",
    "                                  date_created=time.strftime(\"%a, %d %b %Y %H:%M:%S +0000\", time.localtime()),\n",
    "                                  description=\"Conversion of ApolloCar3D dataset into MS-COCO format\")\n",
    "    self.json_file[\"categories\"] = [dict(name='',  # Category name\n",
    "                                         id=1,  # Id of category\n",
    "                                         skeleton=[],  # Skeleton connections (check constants.py)\n",
    "                                         supercategory='',  # Same as category if no supercategory\n",
    "                                         keypoints=[])]  # Keypoint names\n",
    "    self.json_file[\"images\"] = []  # Empty for initialization\n",
    "    self.json_file[\"annotations\"] = []  # Empty for initialization\n",
    "\n",
    "\n",
    "def process_image(json_file):\n",
    "    \"\"\"\n",
    "    Update image field in json file\n",
    "    \"\"\"\n",
    "    # ------------------\n",
    "    # Add here your code\n",
    "    # -------------------\n",
    "    json_file[\"images\"].append({\n",
    "        'coco_url': \"unknown\",\n",
    "        'file_name': '',  # Image name\n",
    "        'id': 0,  # Image id\n",
    "        'license': 1,  # License type\n",
    "        'date_captured': \"unknown\",  \n",
    "        'width': 0,  # Image width (pixels)\n",
    "        'height': 0})  # Image height (pixels)\n",
    "\n",
    "\n",
    "def process_annotation(json_file):\n",
    "    \"\"\"\n",
    "    Process and include in the json file a single annotation (instance) from a given image\n",
    "    \"\"\"\n",
    "    # ------------------\n",
    "    # Add here your code\n",
    "    # -------------------\n",
    "    json_file[\"annotations\"].append({\n",
    "        'image_id': 0,  # Image id\n",
    "        'category_id': 1,  # Id of the category (like car or person)\n",
    "        'iscrowd': 0,  # 1 to mask crowd regions, 0 if the annotation is not a crowd annotation\n",
    "        'id': 0,  # Id of the annotations\n",
    "        'area': 0,  # Bounding box area of the annotation (width*height)\n",
    "        'bbox': [],  # Bounding box  coordinates (x0, y0, width, heigth), where x0, y0 are the left corner\n",
    "        'num_keypoints': 0,  # number of keypoints\n",
    "        'keypoints': [],  # Flattened list of keypoints [x, y, visibility, x, y, visibility, .. ]\n",
    "        'segmentation': []})  # To add a segmentation of the annotation, empty otherwise"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training\n",
    "\n",
    "We have seen all the elements needed to create your own plugin on a custom dataset. To train the dataset, all OpenPifPaf commands are still valid. There are only two differences:\n",
    "1. Specify the dataset name in the training command. In this case, we have called our dataset _apollo_ during the registration phase, therefore we will have `--dataset=apollo`. \n",
    "2. Include the commands we have created in the data module for this specific dataset, for example `--apollo-square-edge` to define the size of the training crops.\n",
    "\n",
    "A training command may look like this:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```sh\n",
    "python3 -m openpifpaf.train --dataset apollo \\\n",
    "--apollo-square-edge=769 \\\n",
    "--basenet=shufflenetv2k16 --lr=0.00002 --momentum=0.95  --b-scale=5.0 \\\n",
    "--epochs=300 --lr-decay 160 260 --lr-decay-epochs=10  --weight-decay=1e-5 \\\n",
    "--weight-decay=1e-5  --val-interval 10 --loader-workers 16 --apollo-upsample 2 \\\n",
    "--apollo-bmin 2 --batch-size 8\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##  Evaluation\n",
    "Evaluation on the COCO metric is supported by pifpaf and a simple evaluation command may look like this:\n",
    "```sh\n",
    "python3 -m openpifpaf.eval --dataset=apollo --checkpoint <path of the model>\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To evaluate on custom metrics, we would need to define a new metric and add it in the list of metrics, inside the data module. In our case, we have a a DataModule class called __ApolloKp__, and its function _**metrics**_ returns the list of metrics to run. Each metric is defined as a class that inherits from openpifpaf.metric.base.Base\n",
    "\n",
    "For more information, please check how we implemented a simple metric for the ApolloCar3D dataset called MeanPixelError, that calculate mean pixel error and detection rate for a given image."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prediction\n",
    "To run your model trained on a different dataset, you simply need to run the standard OpenPifPaf command specifying your model. A prediction command looks like this:\n",
    "```sh\n",
    "python3 -m openpifpaf.predict --checkpoint <model path>\n",
    "```\n",
    "\n",
    "All the command line options are still valid, check them with:\n",
    "```sh\n",
    "python3 -m openpifpaf.predict --help\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Final remarks\n",
    "We hope you'll find this guide useful to create your own plugin. For more information check the guide section for the {doc}`ApolloCar3D plugin <plugins_apollocar3d>`.\n",
    "\n",
    "Please keep us posted on issues you encounter (using the issue section on GitHub) and especially on your successes! We will be more than happy to add your plugin to the list of OpenPifPaf [related projects](https://openpifpaf.github.io/intro.html#related-projects).\n",
    "\n",
    "Finally, if you find OpenPifPaf useful for your research, we would be happy if you {ref}`cite us <citation>`!\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": ""
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
